package org.sz;

import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;

import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.function.Consumer;

/**
 * Redis文档读取器
 * 支持从Redis中读取数据并转换为Document对象
 */
public class RedisDocumentReader implements DocumentReader, AutoCloseable {

    private final RedisResource redisResource;
    private final JedisPool jedisPool;
    private final ExecutorService executorService;

    public RedisDocumentReader(RedisResource redisResource) {
        this.redisResource = redisResource;
        this.jedisPool = createJedisPool();
        this.executorService = Executors.newFixedThreadPool(
            Math.max(2, Runtime.getRuntime().availableProcessors()));
    }

    /**
     * 创建Jedis连接池
     */
    private JedisPool createJedisPool() {
        JedisPoolConfig poolConfig = new JedisPoolConfig();
        poolConfig.setMaxTotal(redisResource.getMaxConnections());
        poolConfig.setMaxIdle(redisResource.getMaxIdle());
        poolConfig.setMinIdle(redisResource.getMinIdle());
        poolConfig.setTestOnBorrow(true);
        poolConfig.setTestOnReturn(true);
        
        if (redisResource.getPassword() != null && !redisResource.getPassword().isEmpty()) {
            return new JedisPool(poolConfig, 
                redisResource.getHost(), 
                redisResource.getPort(),
                redisResource.getConnectionTimeout(),
                redisResource.getPassword());
        } else {
            return new JedisPool(poolConfig, 
                redisResource.getHost(), 
                redisResource.getPort(),
                redisResource.getConnectionTimeout());
        }
    }

    @Override
    public List<Document> get() {
        List<Document> documents = new ArrayList<>();
        
        try (Jedis jedis = jedisPool.getResource()) {
            Set<String> keys = jedis.keys(redisResource.getKeyPattern());
            
            for (String key : keys) {
                Document doc = processKey(jedis, key);
                if (doc != null) {
                    documents.add(doc);
                }
            }
        }
        
        return documents;
    }

    /**
     * 处理单个Redis键
     */
    private Document processKey(Jedis jedis, String key) {
        String type = jedis.type(key);
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("key", key);
        metadata.put("type", type);
        metadata.put("source", "redis");
        
        String content;
        
        switch (type) {
            case "string":
                content = jedis.get(key);
                break;
            case "hash":
                Map<String, String> hash = jedis.hgetAll(key);
                content = hash.toString();
                metadata.putAll(hash);
                break;
            case "list":
                List<String> list = jedis.lrange(key, 0, -1);
                content = list.toString();
                metadata.put("size", list.size());
                break;
            case "set":
                Set<String> set = jedis.smembers(key);
                content = set.toString();
                metadata.put("size", set.size());
                break;
            case "zset":
                Set<String> zset = (Set<String>) jedis.zrange(key, 0, -1);
                content = zset.toString();
                metadata.put("size", zset.size());
                break;
            default:
                content = "Unsupported type: " + type;
        }
        
        return new Document(content, metadata);
    }

    /**
     * 流式处理Redis数据
     */
    public void streamDocuments(Consumer<Document> consumer) {
        try (Jedis jedis = jedisPool.getResource()) {
            Set<String> keys = jedis.keys(redisResource.getKeyPattern());
            
            for (String key : keys) {
                Document doc = processKey(jedis, key);
                if (doc != null) {
                    consumer.accept(doc);
                }
            }
        }
    }

    /**
     * 异步获取文档列表
     */
    public Future<List<Document>> getAsync() {
        return executorService.submit(this::get);
    }

    @Override
    public void close() {
        if (executorService != null && !executorService.isShutdown()) {
            executorService.shutdown();
        }
        if (jedisPool != null && !jedisPool.isClosed()) {
            jedisPool.close();
        }
    }
}